package com.wyp.springbootinit.utils;


import cn.hutool.core.lang.Pair;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * yolo识别工具类
 *
 * @author 4413
 */
public class YoloUtils {
    public static Pair<Integer, String> recognize(File file) throws InterruptedException, IOException {
        System.out.println("yolo识别");
        String condaEnvName = "yolo";
        String pythonScriptPath = "D:\\code\\python\\yolov5-master\\detect.py";

        // 构造conda命令
        List<String> command = new ArrayList<>();
        command.add("conda");
        command.add("run");
        command.add("-n");
        command.add(condaEnvName);
        command.add("python");
        command.add(pythonScriptPath);
        command.add("--source");
        command.add(file.getAbsolutePath());
        command.add("--device");
        command.add("0");

        // 构造ProcessBuilder对象
        ProcessBuilder processBuilder = new ProcessBuilder(command);
        processBuilder.redirectErrorStream(true);

        // 启动进程并等待完成
        Process process = processBuilder.start();

        BufferedReader stdInput = new BufferedReader(new InputStreamReader(process.getInputStream()));

        String s;
        int cnt = 0;
        String result = "";
        while ((s = stdInput.readLine()) != null) {
            if ("image".equals(s.split(" ")[0])) {
                System.out.println(s);
                cnt = Integer.parseInt(s.split(" ")[4]);
            }
            if (s.contains("Results")) {
                System.out.println(s);
                result = s.split("\\\\")[6].split("\u001B")[0];
            }
        }
        process.waitFor();

        // 使用cnt和result构造返回值
        return new Pair<>(cnt, result);
    }
}
